#include "kbest_index.h"
#include "kbest.h"
#include <filesystem>

KBestIndex::KBestIndex(const ConfigParser& config) {
    k_f = config.get_int("k_f", 32);
    efs = config.get_int("efs", 96);
    efc = config.get_int("efc", 300);
    A = config.get_int("A", 1);
    graph_opt_iter = config.get_int("graph_opt_iter", 9);
    metric = config.get_string("metric", "L2");
    init_builder_type = config.get_string("init_builder_type", "RNNDescent");
    index_type = config.get_string("index_type", "TSDG");
    consecutive = config.get_int("consecutive", 20);
    reorder = config.get_int("reorder", 1);
    level = config.get_int("level", 0);
    adding_pref = config.get_int("adding_pref", 23);
    patience = config.get_int("patience", 32);
    NUMA_ENABLED = config.get_bool("NUMA_ENABLED", false);
    num_numa_nodes = config.get_int("num_numa_nodes", 4);
    index_path = config.get_string("index_path", "searcher.kbest");
    save_or_load = config.get_string("save_or_load", "neither");
}


void KBestIndex::build(const Datasets* data) {
    extern bool NUMA_ENABLED;
    extern int num_numa_nodes;
    best_build = std::make_unique<KBest>(
        data->d, k_f, efc, A, graph_opt_iter, 
        metric, init_builder_type, index_type
    );
    if (save_or_load == "load") {
        best_build->Load(index_path.c_str());
    } else {
        best_build->Add(data->nb, data->xb, consecutive, reorder, level);
        best_build->BuildSearcher();
        if (save_or_load == "save") {
			std::filesystem::path dir_path = std::filesystem::path(index_path).parent_path();
			if (!dir_path.empty() && !std::filesystem::exists(dir_path)) {
                std::filesystem::create_directories(dir_path);
            }
            best_build->Save(index_path.c_str());
        }
    }
}

void KBestIndex::prepare_search() {
    best_build->SetEf(efs);
    best_build->SetEarlyStoppingParams(adding_pref, patience);
}

void KBestIndex::search_single(const float* xq, int k, float* distances, int64_t* labels) {
    int n = 1;
    int num_thread = 1;
    best_build->Search(n, xq, k, distances, labels, num_thread);
}

void KBestIndex::search_batch(const int nq, const float* xq, int k, float* distances, int64_t* labels) {
    int num_threads = 80;
    best_build->Search(nq, xq, k, distances, labels, num_threads);
}

void KBestIndex::search_batch(const int nq, const float* xq, int k, float* distances, int64_t* labels, int num_threads) {
    best_build->Search(nq, xq, k, distances, labels, num_threads);
}